#include "avl.h"


void avl_preorder_dump(const avl_node *node, int *dst, int  *pcnt)
{
    if (node == NULL)
        return;
    dst[*pcnt] = node->val;
    (*pcnt)++;

    avl_preorder_dump( node->lchild, dst, pcnt);
    avl_preorder_dump( node->rchild, dst, pcnt);
}

void avl_inorder_dump(const avl_node *node, int *dst, int  *pcnt)
{
    if (node == NULL)
        return;
    avl_inorder_dump( node->lchild, dst, pcnt);
     dst[*pcnt] = node->val;
    (*pcnt)++;
    avl_inorder_dump( node->rchild, dst, pcnt);
}
void avl_postorder_dump(const avl_node *node, int *dst, int  *pcnt)
{
    if (node == NULL)
        return;
    avl_postorder_dump( node->lchild, dst, pcnt);
    avl_postorder_dump( node->rchild, dst, pcnt);
    dst[*pcnt] = node->val;
    (*pcnt)++;
}


void avl_preorder_dump(const avl_node *node, std::vector<int> &echo)
{
    std::stack<const avl_node *> st;
    const avl_node *root = node;

    while ( root || !st.empty()) {

        if ( root ) {
            st.push( root );
            echo.push_back(root->val);
            root = root->lchild;
            continue;
        }

        root = st.top()->rchild;
        st.pop();
    }

}

void avl_inorder_dump(const avl_node *node, std::vector<int> &echo)
{
    const avl_node *root = node;
    std::stack<const avl_node *> st{};

    while (!st.empty() || root != nullptr) {
        if ( root ) {
            st.push(root);
            root = root->lchild;
            continue;
        }
       
        root = st.top()->rchild;
        echo.push_back( st.top()->val );    
        st.pop();
    }
}
void avl_postorder_dump(const avl_node *node, std::vector<int> &echo)
{
    const avl_node *root = node;
    std::stack<const avl_node *> st{};
    const avl_node *pre  = NULL;

    while (!st.empty() || root != nullptr) {
        if ( root ) {
            st.push(root);
            root = root->lchild;
            continue;
        }

        if (st.top()->rchild && pre != st.top()->rchild) {
            root = st.top()->rchild;
        }
        else {
            echo.push_back( st.top()->val );    
            pre = st.top();
            st.pop();
        }
    }
}




int avl_get_height(const avl_node *node) {
    if ( NULL == node)
        return 0;

    int lh = node->lchild == NULL ? 0 : node->lchild->height;
    int rh = node->rchild == NULL ? 0 : node->rchild->height;
    return  lh > rh ? lh + 1 : rh + 1;
}

enum AVL_UB_TYPE avl_get_ub_type(avl_node *node) {
    if (NULL == node)
        return AVL_NORMAL;

    int lh = avl_get_height( node->lchild );
    int rh = avl_get_height( node->rchild );

    if ( abs(lh - rh) <= 1) {
        return AVL_NORMAL;
    }   

    // printf("lh: %d, rh: %d\n", lh, rh);
    if ( lh > rh) {
        int llh = avl_get_height( node->lchild->lchild );
        if ( llh + 1 == lh) {
            return AVL_LLL_UB;
        }
        else {
            return AVL_LR_UB;
        }
    }
    else {
        int rrh = avl_get_height( node->rchild->rchild);

    //    printf("rrh: %d\n", rrh);
        if ( rrh + 1 == rh ) {
            return AVL_RRR_UB;
        }
        else {
            return AVL_RL_UB;
        }
    }
}
avl_node * avl_adjust(avl_node *node, enum AVL_UB_TYPE tp)
{
    if ( node == NULL )
        return NULL;

    avl_node *ret = node;
    switch (tp)
    {
        case AVL_LLL_UB:
            ret = avl_rotateR( node );
            break;
        case AVL_RRR_UB:
            ret = avl_rotateL( node );
            break;
        case AVL_RL_UB:
            ret = avl_rotateRL( node );
            break;
        case AVL_LR_UB:
            ret = avl_rotateLR( node );
            break;
        default:
            break;
    }

    return ret;
}
avl_node * avl_make_new_node(int v) {
    avl_node * ret = (avl_node *) malloc(sizeof(avl_node));
    ret->val = v;
    ret->lchild = NULL;
    ret->rchild = NULL;
    ret->parent = NULL;
    ret->height = 1;
    return ret;
}
/*
 Left Rotate

     p                         p
     |                         |
     x                         y
    / \                       / \
   A   y     --->            x   C
      / \                   / \
     B   C                 A   B

depth(y) - depth(A) = 2
depth(C) - depth(A) = 1
depth(B) = depth(C) or depth(B) + 1 = depth(C) 


after rotate
depth(x) = depth(B) + 1
depth(y) = depth(x) + 1


let p be y's parent, y be p's son
let x be B's parent, B be x's right son
let y be x's parent, x be y's left  son

 */

avl_node * avl_rotateL(avl_node *tar)
{
    avl_node * x = tar;
    avl_node * y = tar->rchild;
    avl_node * p = tar->parent;

    avl_node * A = tar->lchild;
    avl_node * B = y->lchild;
    avl_node * C = y->rchild;

    if ( p != NULL ) {
        if ( p->lchild == x) {
            p->lchild = y;
        }
        else { 
            p->rchild = y;
        }
    }
    y->parent = p;

    x->rchild = B;
    if ( B != NULL )
        B->parent = x;

    y->lchild = x;
    x->parent = y;


    // x->height = 1;
    // if ( B != NULL )
    //     x->height = B->height + 1;
    x->height = avl_get_height( x );


    y->height = x->height + 1;


    return y;
}
/*
 Right Rotate

     p                         p
     |                         |
     x                         y
    / \                       / \
   y   C     --->            A   x
  / \                           / \
 A   B                         B   C

depth(y) - depth(C) = 2
depth(A) - depth(C) = 1
depth(A) = depth(B) or depth(B) + 1 = depth(A) 


after rotate
depth(x) = depth(B) + 1
depth(y) = depth(x) + 1


let p be y's parent, y be p's son
let x be B's parent, B be x's left  son
let y be x's parent, x be y's right son
*/

avl_node * avl_rotateR(avl_node *tar)
{
    avl_node * p = tar->parent;
    avl_node * x = tar;
    avl_node * y = x->lchild;
    avl_node * A = y->lchild;
    avl_node * B = y->rchild;
    avl_node * C = x->rchild;

    y->parent = p;
    if ( p != NULL ) {
        if (p->lchild == x) {
            p->lchild = y;
        }
        else {
            p->rchild = y;
        }
    }

    x->lchild = B;
    if ( B != NULL )
        B->parent = x;

    x->parent = y;
    y->rchild = x;


    x->height = avl_get_height( x );
    // x->height = 1;
    // if ( B != NULL )
    //     x->height = B->height + 1;

    y->height = x->height + 1;

    return y;
}
/*
    Rotate LR

        p                                   p                               p
        |                                   |                               |
        x                                   x                               z
       / \                                 / \                             / \ 
      y   D     RotateL(y)                z   D     Rotate(x)            y     x
     /  \                                / \                            / \   / \
    A    z                              y   C                          A   B C   D 
        / \                            / \
       B   C                          A   B

before rotate:

depth(y) - depth(D) = 2
depth(z) - depth(D) = 1
depth(A) = depth(D)

*/

avl_node * avl_rotateLR(avl_node *tar)
{
    avl_node * x = tar;
    avl_node * y = tar->lchild;
    avl_node * z = y->rchild;

    avl_rotateL(y);
    avl_rotateR(x);

    return z;
}

/*
Rotate RL


          p                             p                              p
          |                             |                              |
          x                             x                              z
         / \                           / \                            / \
        A   y       RotateR(y)        A   z        RotateL(x)        x   y
           / \                           / \                        / \ / \
          z   D                         B   y                      A  B C  D
         / \                               / \
        B   C                             C   D 

before rotate:

depth(y) - depth(A) = 2
depth(z) - depth(A) = 1
depth(A) = depth(D)

*/


avl_node * avl_rotateRL(avl_node *tar)
{
    avl_node * x = tar;
    avl_node * y = tar->rchild;
    avl_node * z = y->lchild;

    avl_rotateR(y);
    avl_rotateL(x);

    return z;
}


bool avl_insert(avl_node **root, int target)
{

    avl_node **pos = root;

    // find pos to insert new val
    while ( *pos ) {

        // if there already have a value
        if ( target == (pos[0]->val)) {
            return false;
        }

        if ( target > (pos[0]->val)) {
            // if can insert in node's rchild
            if (pos[0]->rchild == NULL)
                break;
            pos = &(pos[0]->rchild);
        }
        else {
            // if can insert in node's lchild
            if (pos[0]->lchild == NULL)
                break;
            pos = &(pos[0]->lchild);
        }
    }

    avl_node *node = avl_make_new_node( target );
    if ( *pos == NULL) {
        *pos = node;
    }
    else {
        node->parent = pos[0];
        if (pos[0]->val > target) {
            pos[0]->lchild = node;
        }
        else {
            pos[0]->rchild = node;
        } 
        
        avl_rebalanced( pos[0] );
        avl_update_new_tree_root( root );
    }

   return true;
}




avl_node * avl_find_node(avl_node **root, int target)
{
    if (root == NULL)
        return NULL;

    avl_node *pos = *root;
    while ( pos ) {

        if ( target == pos->val )
            return pos;
        
        if ( target > pos->val ) {
            pos = pos->rchild;
        }
        else {
            pos = pos->lchild;
        }
    }

    return pos;
}


bool avl_find( avl_node **root, int target)
{
    return avl_find_node( root, target) != NULL;
}


bool avl_delete( avl_node **root, int target)
{
    avl_node *delete_node = avl_find_node( root, target);
    if ( delete_node == NULL)
        return false;
    
    avl_node *rep_node = avl_get_rep_node( delete_node );

    if ( rep_node == NULL ) {
        if (delete_node == *root) {
            *root = NULL;
        }
        else {
            avl_node *pa = delete_node->parent;
            avl_unlinked_from_parent( delete_node );
            avl_rebalanced( pa );
        }
        avl_free_node( delete_node );
    }
    else {
        delete_node->val = rep_node->val;
        avl_node *pa = rep_node->parent;
        avl_unlinked_from_parent( rep_node );
        avl_rebalanced( pa );
        avl_free_node( rep_node );
    }

    avl_update_new_tree_root( root );
    return true;
}


avl_node * avl_get_rep_node( avl_node * delete_node )
{   
    avl_node *left_tree_last = delete_node->lchild;
    while ( left_tree_last != NULL && left_tree_last->rchild != NULL) {
        left_tree_last = left_tree_last->rchild;
    }

    if ( left_tree_last != NULL)
        return left_tree_last;
    
    avl_node *right_tree_first = delete_node->rchild;
    while ( right_tree_first != NULL && right_tree_first->lchild != NULL ) {
        right_tree_first = right_tree_first->lchild;
    }

    return right_tree_first;
}


void    avl_free_node(avl_node * node)
{
    if ( NULL == node )
        return ;
    avl_free_node( node->lchild);
    avl_free_node( node->rchild);
    free( node );
}
void   avl_unlinked_from_parent( avl_node *node)
{
    if (node == NULL || node->parent == NULL)
        return;
    
    avl_node * pa = node->parent;
    if ( pa->lchild == node )
        pa->lchild = NULL;
    else 
        pa->rchild = NULL;
    
    node->parent = NULL;
}


void avl_rebalanced(avl_node *node)
{
    if ( node == NULL )
        return ;
    avl_node *cur = node;

    while ( cur != NULL) {


        // printf("[%d]: %d\n", cur->height, cur->val);
        int old_height = cur->height;

        AVL_UB_TYPE tp = avl_get_ub_type( cur );
        cur = avl_adjust( cur, tp);

        // printf("rotate TYPE: %s\n", avl_ub_strs[tp]);

        int new_height = avl_get_height( cur );
        
        if ( old_height == new_height)
            break;



        cur->height = new_height;
        // printf("[%d]: %d\n", cur->height, cur->val);
        
        cur = cur->parent;
    }


    // if ( cur ) {
    //     cur->height = avl_get_height( cur );
    // }

    // printf("cur->val: %d, cur->height: %d\n", cur->val, cur->height);
}


void avl_update_new_tree_root(avl_node ** root)
{
    if (root == NULL )
        return ;
    while ( *root != NULL && root[0]->parent != NULL ) {
        *root = root[0]->parent;
    }
}



bool avl_verify(const avl_node *root)
{

    std::stack<const avl_node *> st;
    const avl_node *pre = NULL;

    while ( root || !st.empty()) {
        if ( root ) {
            st.push( root );
            root = root->lchild;
            continue;
        }
        
        if ( st.top()->rchild && pre != st.top()->rchild) {
            root = st.top()->rchild;
        }
        else {
            
            if (st.top()->height != avl_get_height( st.top())) {
                return false;
            }
            pre = st.top();
            st.pop();
        }
    }

    return true;
}